{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\ProgramData\\Anaconda3\\lib\\site-packages\\h5py\\__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n",
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import keras.backend as K\n",
    "from keras.layers import Input, Lambda\n",
    "from keras.models import Model\n",
    "from keras.optimizers import Adam\n",
    "from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping\n",
    "from nets.yolo3 import yolo_body\n",
    "from nets.loss import yolo_loss\n",
    "from keras.backend.tensorflow_backend import set_session\n",
    "from utils.utils import get_random_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Create YOLOv3 model with 9 anchors and 1 classes.\n",
      "Load weights model_data/yolo_weights.h5.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\ProgramData\\Anaconda3\\lib\\site-packages\\keras\\engine\\topology.py:3473: UserWarning: Skipping loading of weights for layer conv2d_59 due to mismatch in shape ((1, 1, 1024, 18) vs (255, 1024, 1, 1)).\n",
      "  weight_values[i].shape))\n",
      "C:\\ProgramData\\Anaconda3\\lib\\site-packages\\keras\\engine\\topology.py:3473: UserWarning: Skipping loading of weights for layer conv2d_59 due to mismatch in shape ((18,) vs (255,)).\n",
      "  weight_values[i].shape))\n",
      "C:\\ProgramData\\Anaconda3\\lib\\site-packages\\keras\\engine\\topology.py:3473: UserWarning: Skipping loading of weights for layer conv2d_67 due to mismatch in shape ((1, 1, 512, 18) vs (255, 512, 1, 1)).\n",
      "  weight_values[i].shape))\n",
      "C:\\ProgramData\\Anaconda3\\lib\\site-packages\\keras\\engine\\topology.py:3473: UserWarning: Skipping loading of weights for layer conv2d_67 due to mismatch in shape ((18,) vs (255,)).\n",
      "  weight_values[i].shape))\n",
      "C:\\ProgramData\\Anaconda3\\lib\\site-packages\\keras\\engine\\topology.py:3473: UserWarning: Skipping loading of weights for layer conv2d_75 due to mismatch in shape ((1, 1, 256, 18) vs (255, 256, 1, 1)).\n",
      "  weight_values[i].shape))\n",
      "C:\\ProgramData\\Anaconda3\\lib\\site-packages\\keras\\engine\\topology.py:3473: UserWarning: Skipping loading of weights for layer conv2d_75 due to mismatch in shape ((18,) vs (255,)).\n",
      "  weight_values[i].shape))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Freeze the first 249 layers of total 252 layers.\n",
      "load weights\n",
      "Train on 180 samples, val on 20 samples, with batch size 4.\n",
      "Epoch 51/85\n",
      "45/45 [==============================] - 45s 998ms/step - loss: 13.8610 - val_loss: 17.4213\n",
      "Epoch 52/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 13.5634 - val_loss: 16.5080\n",
      "Epoch 53/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 13.9860 - val_loss: 18.2703\n",
      "Epoch 54/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 13.6606 - val_loss: 19.4278\n",
      "Epoch 55/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 13.6437 - val_loss: 16.1858\n",
      "Epoch 56/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 13.4429 - val_loss: 16.8055\n",
      "Epoch 57/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 13.6745 - val_loss: 17.8227\n",
      "Epoch 58/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 13.6393 - val_loss: 18.6143\n",
      "Epoch 59/85\n",
      "45/45 [==============================] - 30s 674ms/step - loss: 13.1447 - val_loss: 17.5637\n",
      "Epoch 60/85\n",
      "45/45 [==============================] - 30s 669ms/step - loss: 13.1005 - val_loss: 17.0495\n",
      "Epoch 61/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 13.2372 - val_loss: 18.3155\n",
      "Epoch 62/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 13.2527 - val_loss: 17.5826\n",
      "Epoch 63/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 13.3895 - val_loss: 16.4717\n",
      "Epoch 64/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 13.5608 - val_loss: 16.5224\n",
      "Epoch 65/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 12.9858 - val_loss: 15.6877\n",
      "Epoch 66/85\n",
      "45/45 [==============================] - 30s 669ms/step - loss: 13.3306 - val_loss: 15.9866\n",
      "Epoch 67/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 12.9819 - val_loss: 16.7652\n",
      "Epoch 68/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 12.8585 - val_loss: 14.1420\n",
      "Epoch 69/85\n",
      "45/45 [==============================] - 30s 669ms/step - loss: 13.0023 - val_loss: 16.6869\n",
      "Epoch 70/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 12.7901 - val_loss: 17.6099\n",
      "Epoch 71/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 12.7061 - val_loss: 17.7580\n",
      "Epoch 72/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 12.6000 - val_loss: 17.9031\n",
      "Epoch 73/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 12.9149 - val_loss: 16.8893\n",
      "Epoch 74/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 12.2931 - val_loss: 16.1376\n",
      "Epoch 75/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 12.5736 - val_loss: 16.4398\n",
      "Epoch 76/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 12.4051 - val_loss: 18.8736\n",
      "Epoch 77/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 12.5040 - val_loss: 16.8844\n",
      "Epoch 78/85\n",
      "45/45 [==============================] - 30s 666ms/step - loss: 12.7062 - val_loss: 15.9262\n",
      "Epoch 79/85\n",
      "45/45 [==============================] - 30s 666ms/step - loss: 12.4012 - val_loss: 15.7031\n",
      "Epoch 80/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 12.5055 - val_loss: 16.7820\n",
      "Epoch 81/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 12.2745 - val_loss: 17.3973\n",
      "Epoch 82/85\n",
      "45/45 [==============================] - 30s 667ms/step - loss: 12.9200 - val_loss: 16.2234\n",
      "Epoch 83/85\n",
      "45/45 [==============================] - 30s 666ms/step - loss: 12.2376 - val_loss: 16.5923\n",
      "Epoch 84/85\n",
      "45/45 [==============================] - 30s 666ms/step - loss: 12.3284 - val_loss: 16.6936\n",
      "Epoch 85/85\n",
      "45/45 [==============================] - 30s 668ms/step - loss: 12.7946 - val_loss: 16.2638\n"
     ]
    }
   ],
   "source": [
    "#---------------------------------------------------#\n",
    "#   获得类和先验框\n",
    "#---------------------------------------------------#\n",
    "def get_classes(classes_path):\n",
    "    '''loads the classes'''\n",
    "    with open(classes_path) as f:\n",
    "        class_names = f.readlines()\n",
    "    class_names = [c.strip() for c in class_names]\n",
    "    return class_names\n",
    "\n",
    "def get_anchors(anchors_path):\n",
    "    '''loads the anchors from a file'''\n",
    "    with open(anchors_path) as f:\n",
    "        anchors = f.readline()\n",
    "    anchors = [float(x) for x in anchors.split(',')]\n",
    "    return np.array(anchors).reshape(-1, 2)\n",
    "\n",
    "#---------------------------------------------------#\n",
    "#   训练数据生成器\n",
    "#---------------------------------------------------#\n",
    "def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):\n",
    "    '''data generator for fit_generator'''\n",
    "    n = len(annotation_lines)\n",
    "    i = 0\n",
    "    while True:\n",
    "        image_data = []\n",
    "        box_data = []\n",
    "        for b in range(batch_size):\n",
    "            if i==0:\n",
    "                np.random.shuffle(annotation_lines)\n",
    "            image, box = get_random_data(annotation_lines[i], input_shape, random=True)\n",
    "            image_data.append(image)\n",
    "            box_data.append(box)\n",
    "            i = (i+1) % n\n",
    "        image_data = np.array(image_data)\n",
    "        box_data = np.array(box_data)\n",
    "        y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)\n",
    "        yield [image_data, *y_true], np.zeros(batch_size)\n",
    "\n",
    "\n",
    "#---------------------------------------------------#\n",
    "#   读入xml文件，并输出y_true\n",
    "#---------------------------------------------------#\n",
    "def preprocess_true_boxes(true_boxes, input_shape, anchors, num_classes):\n",
    "\n",
    "    assert (true_boxes[..., 4]<num_classes).all(), 'class id must be less than num_classes'\n",
    "    # 一共有三个特征层数\n",
    "    num_layers = len(anchors)//3\n",
    "    # 先验框\n",
    "    # 678为116,90,  156,198,  373,326\n",
    "    # 345为30,61,  62,45,  59,119\n",
    "    # 012为10,13,  16,30,  33,23,  \n",
    "    anchor_mask = [[6,7,8], [3,4,5], [0,1,2]] if num_layers==3 else [[3,4,5], [1,2,3]]\n",
    "\n",
    "    true_boxes = np.array(true_boxes, dtype='float32')\n",
    "    input_shape = np.array(input_shape, dtype='int32') # 416,416\n",
    "    # 读出xy轴，读出长宽\n",
    "    # 中心点(m,n,2)\n",
    "    boxes_xy = (true_boxes[..., 0:2] + true_boxes[..., 2:4]) // 2\n",
    "    boxes_wh = true_boxes[..., 2:4] - true_boxes[..., 0:2]\n",
    "    # 计算比例\n",
    "    true_boxes[..., 0:2] = boxes_xy/input_shape[:]\n",
    "    true_boxes[..., 2:4] = boxes_wh/input_shape[:]\n",
    "\n",
    "    # m张图\n",
    "    m = true_boxes.shape[0]\n",
    "    # 得到网格的shape为13,13;26,26;52,52\n",
    "    grid_shapes = [input_shape//{0:32, 1:16, 2:8}[l] for l in range(num_layers)]\n",
    "    # y_true的格式为(m,13,13,3,85)(m,26,26,3,85)(m,52,52,3,85)\n",
    "    y_true = [np.zeros((m,grid_shapes[l][0],grid_shapes[l][1],len(anchor_mask[l]),5+num_classes),\n",
    "        dtype='float32') for l in range(num_layers)]\n",
    "    # [1,9,2]\n",
    "    anchors = np.expand_dims(anchors, 0)\n",
    "    anchor_maxes = anchors / 2.\n",
    "    anchor_mins = -anchor_maxes\n",
    "    # 长宽要大于0才有效\n",
    "    valid_mask = boxes_wh[..., 0]>0\n",
    "\n",
    "    for b in range(m):\n",
    "        # 对每一张图进行处理\n",
    "        wh = boxes_wh[b, valid_mask[b]]\n",
    "        if len(wh)==0: continue\n",
    "        # [n,1,2]\n",
    "        wh = np.expand_dims(wh, -2)\n",
    "        box_maxes = wh / 2.\n",
    "        box_mins = -box_maxes\n",
    "\n",
    "        # 计算真实框和哪个先验框最契合\n",
    "        intersect_mins = np.maximum(box_mins, anchor_mins)\n",
    "        intersect_maxes = np.minimum(box_maxes, anchor_maxes)\n",
    "        intersect_wh = np.maximum(intersect_maxes - intersect_mins, 0.)\n",
    "        intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]\n",
    "        box_area = wh[..., 0] * wh[..., 1]\n",
    "        anchor_area = anchors[..., 0] * anchors[..., 1]\n",
    "        iou = intersect_area / (box_area + anchor_area - intersect_area)\n",
    "        # 维度是(n) 感谢 消尽不死鸟 的提醒\n",
    "        best_anchor = np.argmax(iou, axis=-1)\n",
    "\n",
    "        for t, n in enumerate(best_anchor):\n",
    "            for l in range(num_layers):\n",
    "                if n in anchor_mask[l]:\n",
    "                    # floor用于向下取整\n",
    "                    i = np.floor(true_boxes[b,t,0]*grid_shapes[l][1]).astype('int32')\n",
    "                    j = np.floor(true_boxes[b,t,1]*grid_shapes[l][0]).astype('int32')\n",
    "                    # 找到真实框在特征层l中第b副图像对应的位置\n",
    "                    k = anchor_mask[l].index(n)\n",
    "                    c = true_boxes[b,t, 4].astype('int32')\n",
    "                    y_true[l][b, j, i, k, 0:4] = true_boxes[b,t, 0:4]\n",
    "                    y_true[l][b, j, i, k, 4] = 1\n",
    "                    y_true[l][b, j, i, k, 5+c] = 1\n",
    "\n",
    "    return y_true\n",
    "\n",
    "\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allocator_type = 'BFC' #A \"Best-fit with coalescing\" algorithm, simplified from a version of dlmalloc.\n",
    "config.gpu_options.per_process_gpu_memory_fraction = 0.7\n",
    "config.gpu_options.allow_growth = True\n",
    "set_session(tf.Session(config=config)) \n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # 标签的位置\n",
    "    annotation_path = '2020_3_train.txt'\n",
    "    # 获取classes和anchor的位置\n",
    "    classes_path = 'model_data/voc_classes.txt'    \n",
    "    anchors_path = 'model_data/yolo_anchors.txt'\n",
    "    # 预训练模型的位置\n",
    "    weights_path = 'model_data/yolo_weights.h5'  #预训练模型\n",
    "    # 获得classes和anchor\n",
    "    class_names = get_classes(classes_path)\n",
    "    anchors = get_anchors(anchors_path)\n",
    "    # 一共有多少类\n",
    "    num_classes = len(class_names)\n",
    "    num_anchors = len(anchors)\n",
    "    # 训练后的模型保存的位置\n",
    "    log_dir = 'logs/'\n",
    "    # 输入的shape大小\n",
    "    input_shape = (416,416)\n",
    "\n",
    "    # 清除session\n",
    "    K.clear_session()\n",
    "\n",
    "    # 输入的图像为\n",
    "    image_input = Input(shape=(None, None, 3))\n",
    "    h, w = input_shape\n",
    "\n",
    "    # 创建yolo模型\n",
    "    print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))\n",
    "    model_body = yolo_body(image_input, num_anchors//3, num_classes)\n",
    "    \n",
    "    # 载入预训练权重\n",
    "    print('Load weights {}.'.format(weights_path))\n",
    "    model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)\n",
    "    \n",
    "    # y_true为13,13,3,85\n",
    "    # 26,26,3,85\n",
    "    # 52,52,3,85\n",
    "    y_true = [Input(shape=(h//{0:32, 1:16, 2:8}[l], w//{0:32, 1:16, 2:8}[l], \\\n",
    "        num_anchors//3, num_classes+5)) for l in range(3)]\n",
    "\n",
    "    # 输入为*model_body.input, *y_true\n",
    "    # 输出为model_loss\n",
    "    loss_input = [*model_body.output, *y_true]\n",
    "    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',\n",
    "        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(loss_input)\n",
    "\n",
    "    model = Model([model_body.input, *y_true], model_loss)\n",
    "\n",
    "    freeze_layers = 249\n",
    "    for i in range(freeze_layers): model_body.layers[i].trainable = False\n",
    "    print('Freeze the first {} layers of total {} layers.'.format(freeze_layers, len(model_body.layers)))\n",
    "\n",
    "    # 训练参数设置\n",
    "    logging = TensorBoard(log_dir=log_dir)\n",
    "    checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',\n",
    "        monitor='val_loss', save_weights_only=True, save_best_only=False, period=5)\n",
    "    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=2, verbose=1)\n",
    "    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=6, verbose=1)\n",
    "\n",
    "    # 0.1用于验证，0.9用于训练\n",
    "    val_split = 0.1\n",
    "    with open(annotation_path) as f:\n",
    "        lines = f.readlines()\n",
    "    np.random.seed(10101)\n",
    "    np.random.shuffle(lines)\n",
    "    np.random.seed(None)\n",
    "    num_val = int(len(lines)*val_split)\n",
    "    num_train = len(lines) - num_val\n",
    "\n",
    "    #微调，训练后面几层\n",
    "    # 调整非主干模型first\n",
    "#     if True:\n",
    "#         model.compile(optimizer=Adam(lr=1e-3), loss={\n",
    "#             'yolo_loss': lambda y_true, y_pred: y_pred})\n",
    "\n",
    "#         batch_size = 16\n",
    "#         print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))\n",
    "#         model.fit_generator(data_generator(lines[:num_train], batch_size, input_shape, anchors, num_classes),\n",
    "#                 steps_per_epoch=max(1, num_train//batch_size),\n",
    "#                 validation_data=data_generator(lines[num_train:], batch_size, input_shape, anchors, num_classes),\n",
    "#                 validation_steps=max(1, num_val//batch_size),\n",
    "#                 epochs=10,\n",
    "#                 initial_epoch=0,\n",
    "#                 callbacks=[logging, checkpoint])\n",
    "#         model.save_weights(log_dir + 'trained_weights_stage_2_10.h5')\n",
    "\n",
    "    #在之前的模型上继续训练,此时注释掉非主干模型\n",
    "    print('load weights')\n",
    "    model.load_weights('D:/yolo3-keras-master/logs/ep050-loss13.956-val_loss17.201.h5')\n",
    "    \n",
    "    \n",
    "    for i in range(freeze_layers): model_body.layers[i].trainable = True\n",
    "\n",
    "    # 解冻后训练\n",
    "    if True:\n",
    "        model.compile(optimizer=Adam(lr=1e-4), loss={\n",
    "            'yolo_loss': lambda y_true, y_pred: y_pred})\n",
    "\n",
    "        batch_size = 4\n",
    "        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))\n",
    "        model.fit_generator(data_generator(lines[:num_train], batch_size, input_shape, anchors, num_classes),\n",
    "                steps_per_epoch=max(1, num_train//batch_size),\n",
    "                validation_data=data_generator(lines[num_train:], batch_size, input_shape, anchors, num_classes),\n",
    "                validation_steps=max(1, num_val//batch_size),\n",
    "                epochs=85,\n",
    "                initial_epoch=50,\n",
    "                callbacks=[logging, checkpoint])\n",
    "        model.save_weights(log_dir + 'last2_85.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
